/**
 *
 *  plugin_KeyManagerPlugin.cc
 *
 */

#include "plugin_KeyManagerPlugin.h"

#include <fstream>
#include <vector>
#include <trantor/utils/Logger.h>


using namespace drogon;

/**
 * 构造函数
 */
KeyManagerPlugin::KeyManagerPlugin() = default;

/**
 * 初始化并启动插件
 * @param config 配置文件参数
 */
void KeyManagerPlugin::initAndStart(const Json::Value &config) {
    // RSA密钥路径
    std::string privateKeyPath = config["rsa_private"].asString();
    // RSA公钥路径
    std::string publicKeyPath = config["rsa_public"].asString();

    try {
        // 加载RSA密钥
        loadKey(privateKeyPath, true);
        // 加载RSA公钥
        loadKey(publicKeyPath, false);
    } catch (CryptoPP::Exception &e) {
        // 加载RSA密钥失败
        LOG_ERROR << "Failed to load RSA keys: " << e.what();
        // 退出程序
        exit(EXIT_FAILURE);
    }
}

/**
 * 关闭插件
 */
void KeyManagerPlugin::shutdown() {
    /// Shutdown the plugin
}

/**
 * RSA加密方法
 * @param str 明文字符串
 * @return 加密后的字符串
 */
[[maybe_unused]] std::string KeyManagerPlugin::rsaEncrypt(const std::string &str) {
    // 加密后的字符串
    std::string cipher;
    // 随机数生成器
    CryptoPP::AutoSeededRandomPool rng;

    try {
        // RSA加密
        CryptoPP::RSAES_OAEP_SHA_Encryptor e(publicKey);
        // 加密
        CryptoPP::StringSource ss(
                str,
                true,
                new CryptoPP::PK_EncryptorFilter(rng, e, new CryptoPP::StringSink(cipher)
                )
        );
    } catch (CryptoPP::Exception &e) {
        // 加密失败
        LOG_ERROR << "RSA encryption failed: " << e.what();
        // 返回空字符串
        return "";
    }
    // 返回加密后的字符串
    return cipher;
}

/**
 * RSA解密方法
 * @param str 加密后的字符串
 * @erturn 明文字符串
 */
[[maybe_unused]] std::string KeyManagerPlugin::rsaDecrypt(const std::string &str) {
    // 明文字符串
    std::string recovered;
    // 随机数生成器
    CryptoPP::AutoSeededRandomPool rng;

    try {
        // RSA解密
        CryptoPP::RSAES_OAEP_SHA_Decryptor d(privateKey);
        // 解密
        CryptoPP::StringSource ss(
                str,
                true,
                new CryptoPP::PK_DecryptorFilter(rng, d, new CryptoPP::StringSink(recovered)
                )
        );
    } catch (CryptoPP::Exception &e) {
        // 解密失败
        LOG_ERROR << "RSA decryption failed: " << e.what();
        // 返回空字符串
        return "";
    }
    // 返回明文字符串
    return recovered;
}

/**
 * 加载密钥
 * @param keyPath 密钥路径
 * @param isPrivateKey 是否私钥
 */
void KeyManagerPlugin::loadKey(const std::string &keyPath, bool isPrivateKey) {
    // 字节队列
    CryptoPP::ByteQueue bytes;
    // 读取密钥
    CryptoPP::FileSource file(keyPath.c_str(), true, new CryptoPP::Base64Decoder);
    // 传输到字节队列
    file.TransferTo(bytes);
    // 结束传输
    bytes.MessageEnd();
    // 如果加载的是私钥
    if (isPrivateKey) {
        // 加载私钥
        privateKey.Load(bytes);
    } else {
        // 加载公钥
        publicKey.Load(bytes);
    }
}
